#=====================================
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation


def update_fig(num, bars, line, point, labels, v_line, bardata, m_xs, m_ys):
    for i, b in enumerate(bars):
        b.set_height(bardata[num, i])
    line.set_data(np.array([m_xs[num], m_ys[num]]))
    point.set_offsets(np.array([m_xs[num], m_ys[num]]).T)
    for i, text in enumerate(labels):
        text.set_position((m_xs[num][i], m_ys[num][i]))
    v_line.set_alpha(num / m_xs.shape[0])


amps = np.load('amps.npy')
n_qubits = 5
labels = [bin(i)[2:].zfill(n_qubits) for i in range(2**n_qubits)]
fig = plt.figure(figsize=(18, 6))
ax = fig.add_subplot()
bars = plt.bar(labels, amps[0])
plt.xticks(rotation=45)
ax.set_ylim(0, 0.4)
plt.title('Max-Cut')
plt.xlabel('Bit string')
plt.ylabel('Amplitude')
m_x = np.array(
    [14, 16, 16 + np.sqrt(2), 16 + 2 * np.sqrt(2), 16 + np.sqrt(2), 16])
m_y = np.array([
    0.3, 0.3, 0.3 * (1 + 0.1 * np.sqrt(2)), 0.3, 0.3 * (1 - 0.1 * np.sqrt(2)),
    0.3
])
m_x -= 1
m_y += 0.03
f_x = np.array([17, 15, 17, 15, 17, 15])
f_y = np.linspace(0.25, 0.10, 5)
f_y = np.append(f_y, f_y[1])
m_xs = np.linspace(m_x, f_x, len(amps))
m_ys = np.linspace(m_y, f_y, len(amps))

line = plt.plot(m_x, m_y, 'k')[0]
v_line = plt.plot([16, 16], [0.25, 0.1], 'r--', linewidth=4, alpha=0)[0]
point = plt.scatter(m_x, m_y)
labels = []
for i, label in enumerate([4, 0, 1, 2, 3]):
    labels.append(plt.text(m_x[i], m_y[i], label, fontsize=18))

anim = animation.FuncAnimation(fig,
                               update_fig,
                               len(amps),
                               fargs=(bars, line, point, labels, v_line, amps,
                                      m_xs, m_ys),
                               interval=50)
# plt.show()
anim.save('max_cut.gif')
plt.show()